import asyncio
import socket
import logging
import sys
from command import CMD
from worker import HANDLERS, WORKER
from config import config
import numpy as np

logger = logging.getLogger()

class tcp_server_protocol(asyncio.Protocol):
    def __init__(self, logger, connected:asyncio.Event, disconnected:asyncio.Event, keep_first = True) -> None:
        self.connected = connected
        self.disconnected = disconnected
        self.__conn_num = 0
        self.connected.clear()
        self.disconnected.clear()
        self.logger = logger
        self.keep_first = keep_first    # if TRUE, keep fisrt client and refuse followers. Else keep last and close former.
        self.transport = None
        self.write_err = 0

    def connection_made(self, transport:asyncio.Transport):
        if self.keep_first:
            if self.__conn_num == 0:
                self.transport = transport
                self.disconnected.clear()
                self.connected.set()
                self.logger.info(f'Server : TCP Connected from {transport.get_extra_info("peername")}')
            else:
                self.logger.info(f'Server : Another TCP Connected from {transport.get_extra_info("peername")}, close directly.')
                transport.close()
        else:
            if self.transport:
                self.transport.close()
                self.logger.info(f'Server : Another TCP Connected from {transport.get_extra_info("peername")}, close original.')
            self.transport = transport
            self.logger.info(f'Server : TCP Connected from {transport.get_extra_info("peername")}')
            self.disconnected.clear()
            self.connected.set()

        self.__conn_num += 1


    def connection_lost(self, exc):
        self.__conn_num -= 1
        if self.__conn_num == 0:
            self.connected.clear()
            self.disconnected.set()
            if self.transport is not None:
                self.logger.info(f"Server : TCP Connection from {self.transport.get_extra_info('peername')} closed. Error : {exc}. ")


    def return_self(self):
        return self


class TCP_RESULT_SERVER():

    def __init__(self, cfg:config) -> None:
        super().__init__()
        self.cfg = cfg 
        # self._in_q = in_queue
        self._connected = asyncio.Event()
        self._disconnected = asyncio.Event()

        self.protocol = tcp_server_protocol(logger, self._connected, self._disconnected)
        self.transport = self.protocol.transport

        self._loop = asyncio.get_event_loop()

    
    async def launch_server(self):
        try:
            self.tcp_server = await self._loop.create_server(
                self.protocol.return_self, 
                # lambda: tcp_server_protocol(connected),
                family=socket.AF_INET,
                host='0.0.0.0',
                port= self.cfg.tcp_port, 
                start_serving=False)
            
            async with self.tcp_server:
                asyncio.ensure_future(self.tcp_server.serve_forever())
                while True:
                    await self._connected.wait()
                    self.transport = self.protocol.transport

                    await self._disconnected.wait()

        except BaseException as e:
            logger.exception(e)


    def write_view(self, data_view:memoryview):
        try:
            if self.transport is not None and not self.transport.is_closing():
                # logger.debug("Send data len %i, data : %s ..." % (data_view.nbytes, data_view[:20].tolist()))
                self.transport.write(data_view)
        except OSError as e:
            self.write_err += 1
            logger.warning("Send error occured. Total error number %i" % self.write_err)


    def write(self, data:np.ndarray):
        try:
            if self.transport is not None and not self.transport.is_closing():
                # logger.debug("Send data len %i, data : %s ..." % (data_view.nbytes, data_view[:20].tolist()))
                self.transport.write(data.tobytes())
        except OSError as e:
            self.write_err += 1
            logger.warning("Send error occured. Total error number %i" % self.write_err)


    def write_bytes(self, data:bytes):
        try:
            if self.transport is not None and not self.transport.is_closing():
                # logger.debug("Send data len %i, data : %s ..." % (data_view.nbytes, data_view[:20].tolist()))
                self.transport.write(data)
        except OSError as e:
            self.write_err += 1
            logger.warning("Send error occured. Total error number %i" % self.write_err)
    # async def __periodic_send(self):
    #     err = 0
    #     try:
    #         while self.protocol.transport is not None:
    #             data_view = await self._in_q.get()
    #             logger.debug("Send data len %i, data : %s ..." % (data_view.nbytes, data_view[:20].tolist()))
    #             self.protocol.transport.write(data_view)
        
    #     except asyncio.CancelledError as e:
    #         logger.debug("Send task cancelled.")

    #     except OSError as e:
    #         err += 1
    #         logger.warning("Send error occured. Total error number %i" % err)